package com.aiwiown.snackmq.storage.service;

import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.protocol.TopicPartition;
import com.aiwiown.snackmq.storage.StorageService;
import com.aiwiown.snackmq.storage.exception.PartitionNotFoundException;
import com.aiwiown.snackmq.storage.partition.LogPartition;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.stream.Collectors;

/**
 * 传统的分区存储服务实现。
 */
@Slf4j
public class PartitionedStorageService implements StorageService {

    private final Map<TopicPartition, LogPartition> sharedPartitions;
    private final String storagePath;
    // 【修改】: 将此常量作为默认值，而不是硬性限制
    private static final int DEFAULT_MAX_MESSAGES_PER_PULL = 32;

    // 【新增】恢复状态保护机制
    private final Map<TopicPartition, Boolean> partitionRecoveryStatus = new ConcurrentHashMap<>();

    public PartitionedStorageService(String storagePath, Map<TopicPartition, LogPartition> sharedPartitions) {
        this.storagePath = storagePath;
        this.sharedPartitions = sharedPartitions;
    }

    @Override
    public void start() {
        log.info("Partitioned storage service started and is using shared partitions.");
    }

    @Override
    public void close() throws IOException {
        log.info("Partitioned storage service closed.");
    }

    /**
     * 【新增】设置分区的恢复状态
     *
     * @param partition 分区
     * @param isRecovering 是否正在恢复
     */
    public void setPartitionRecoveryStatus(TopicPartition partition, boolean isRecovering) {
        partitionRecoveryStatus.put(partition, isRecovering);
        log.info("分区 {} 恢复状态设置为: {}", partition, isRecovering ? "恢复中" : "已完成");
    }

    /**
     * 【新增】检查分区是否正在恢复中
     *
     * @param partition 分区
     * @return 是否正在恢复中
     */
    public boolean isPartitionRecovering(TopicPartition partition) {
        // 【核心修复】: 默认值应为 true。
        // 在一个分区首次被创建，其状态尚未被记录到 map 中时，
        // 我们必须安全地假定它仍在恢复中，直到它明确通知恢复完成。
        // 之前的默认值 false 是导致竞态条件和数据丢失的根源。
        return partitionRecoveryStatus.getOrDefault(partition, true);
    }

    /**
     * 【最终修复】: 新增此方法以支持并发测试和原子位点生成。
     * 此方法将位点生成的责任委托给 LogPartition，是并发写入的首选方法。
     *
     * @param partition 目标分区
     * @param message   要追加的消息
     * @return 分配给消息的逻辑位点
     * @throws IOException 如果写入失败
     */
    public long append(TopicPartition partition, Message message) throws IOException {
        if (isPartitionRecovering(partition)) {
            log.warn("分区 {} 正在恢复中，跳过单条消息追加操作，返回-1", partition);
            return -1L;
        }

        LogPartition logPartition = sharedPartitions.get(partition);
        if (logPartition != null) {
            // 调用 LogPartition 中那个返回 long 的新 append 方法
            return logPartition.append(message);
        } else {
            log.warn("Attempted to append to a non-existent or inactive partition: {}", partition);
            throw new PartitionNotFoundException("Partition " + partition + " is not active on this broker.");
        }
    }

    @Override
    @Deprecated
    public void append(TopicPartition partition, Message message, long offset) throws IOException {
        // 【新增】检查分区是否正在恢复中
        if (isPartitionRecovering(partition)) {
            log.warn("分区 {} 正在恢复中，跳过单条消息追加操作", partition);
            return;
        }

        LogPartition logPartition = sharedPartitions.get(partition);
        if (logPartition != null) {
            logPartition.append(message, offset);
        } else {
            log.warn("Attempted to append to a non-existent or inactive partition: {}", partition);
        }
    }
    /**
     * 【新增】批量追加消息到指定分区。
     * 这是为了优化状态机应用日志的性能，避免逐条追加。
     *
     * @param partition    目标分区
     * @param messages     要追加的消息列表
     * @param startOffset  此批次消息的起始逻辑位点
     * @param commitLogOffset 此批次在 Raft 日志中的物理位点 (physical offset)
     * @throws IOException 如果写入失败
     */
    @Override
    public void appendBatch(TopicPartition partition, List<Message> messages, long startOffset, long commitLogOffset) throws IOException {
        // 【新增】检查分区是否正在恢复中
        if (isPartitionRecovering(partition)) {
            log.warn("分区 {} 正在恢复中，跳过批量消息追加操作", partition);
            return;
        }

        LogPartition logPartition = sharedPartitions.get(partition);
        if (logPartition != null) {
            logPartition.appendBatch(messages, startOffset, commitLogOffset);
        } else {
            log.warn("Attempted to batch append to a non-existent or inactive partition: {}", partition);
        }
    }

    @Override
    public long getNextOffset(String topic, int partitionId) throws IOException {
        TopicPartition topicPartition = new TopicPartition(topic, partitionId);
        LogPartition logPartition = sharedPartitions.get(topicPartition);
        if (logPartition != null) {
            return logPartition.getNextOffset();
        }
        return 0L;
    }

    public Message read(TopicPartition partition, long offset) throws IOException {
        LogPartition logPartition = sharedPartitions.get(partition);
        if (logPartition != null) {
            return logPartition.read(offset);
        }
        log.warn("Attempted to read from a non-existent or inactive partition: {}", partition);
        return null;
    }

    /**
     * 【核心修复】: 重载 readBatch 方法，使其接受一个 maxMessages 参数。
     * 旧的无参方法现在委托给新的方法，并使用默认值，以保持向后兼容。
     */
    public List<Message> readBatch(TopicPartition tp, long offset) throws IOException {
        return readBatch(tp, offset, DEFAULT_MAX_MESSAGES_PER_PULL);
    }

    public List<Message> readBatch(TopicPartition tp, long offset, int maxMessages) throws IOException {
        LogPartition partition = sharedPartitions.get(tp);
        if (partition == null) {
            throw new PartitionNotFoundException("Partition " + tp + " is not active on this broker.");
        }
        List<Message> messages = new ArrayList<>(maxMessages);
        long currentOffset = offset;
        for (int i = 0; i < maxMessages; i++) {
            Message message = partition.read(currentOffset);
            if (message == null) {
                break;
            }
            messages.add(message);
            currentOffset++;
        }
        return messages;
    }

    /**
     * 【新增】获取指定分区的LogPartition实例。
     * 主要用于诊断和调试目的。
     *
     * @param tp 主题分区
     * @return LogPartition实例，如果分区不存在则返回null
     */
    public LogPartition getLogPartition(TopicPartition tp) {
        return sharedPartitions.get(tp);
    }

    /**
     * 【新增】修复指定分区的损坏段。
     *
     * @param tp 主题分区
     * @throws IOException 如果修复失败
     * @throws PartitionNotFoundException 如果分区不存在
     */
    public void repairPartition(TopicPartition tp) throws IOException {
        LogPartition partition = sharedPartitions.get(tp);
        if (partition == null) {
            throw new PartitionNotFoundException("Partition " + tp + " is not active on this broker.");
        }
        partition.repairCorruptedSegments();
    }

    /**
     * 【新增】获取指定分区中损坏的段列表。
     *
     * @param tp 主题分区
     * @return 损坏的段列表
     * @throws PartitionNotFoundException 如果分区不存在
     */
    public List<Long> getCorruptedSegments(TopicPartition tp) throws PartitionNotFoundException {
        LogPartition partition = sharedPartitions.get(tp);
        if (partition == null) {
            throw new PartitionNotFoundException("Partition " + tp + " is not active on this broker.");
        }
        return partition.getCorruptedSegments();
    }

    /**
     * 【新增】手动修复指定分区的段结构。
     *
     * @param tp 主题分区
     * @throws IOException 如果修复失败
     * @throws PartitionNotFoundException 如果分区不存在
     */
    public void manualRepairPartition(TopicPartition tp) throws IOException {
        LogPartition partition = sharedPartitions.get(tp);
        if (partition == null) {
            throw new PartitionNotFoundException("Partition " + tp + " is not active on this broker.");
        }
        partition.manualRepairSegments();
    }

    /**
     * 【新增】删除指定 Topic 的所有相关物理存储文件。
     * 此方法会扫描存储根目录，删除所有以 "topicName-" 开头的分区目录。
     *
     * @param topicName 要删除的 Topic 名称。
     */
    @Override
    public void deleteTopic(String topicName) {
        log.info("Starting data cleanup for topic '{}' in PartitionedStorageService...", topicName);
        // --- 核心修复：在删除物理文件前，先关闭并移除内存中的 LogPartition 对象 ---
        // 1. 查找并收集所有属于该 Topic 的分区
        List<TopicPartition> partitionsToRemove = this.sharedPartitions.keySet().stream()
                .filter(tp -> tp.getTopic().equals(topicName))
                .collect(Collectors.toList());
        // 2. 遍历并处理这些分区
        for (TopicPartition tp : partitionsToRemove) {
            // 从共享 Map 中移除，这样就不会再有新的操作指向它
            LogPartition partition = this.sharedPartitions.remove(tp);
            if (partition != null) {
                try {
                    // 关闭分区，释放文件句柄
                    partition.close();
                    log.info("Closed and removed LogPartition from memory: {}", tp);
                } catch (IOException e) {
                    log.error("Error closing LogPartition {} during topic deletion. File deletion might fail.", tp, e);
                }
            }
        }
        // 3. 现在文件句柄已释放，可以安全地删除物理文件
        log.info("Deleting all storage files for topic '{}' from path '{}'...", topicName, this.storagePath);
        File storeDir = new File(this.storagePath);
        if (!storeDir.exists() || !storeDir.isDirectory()) {
            log.warn("Storage directory '{}' does not exist or is not a directory. Skipping topic file deletion.", this.storagePath);
            return;
        }
        // 目录名约定为 "topicName-partitionId"
        File[] topicDirs = storeDir.listFiles((dir, name) -> name.startsWith(topicName + "-") && new File(dir, name).isDirectory());

        if (topicDirs != null) {
            for (File topicDir : topicDirs) {
                try {
                    log.info("Deleting partition directory: {}", topicDir.getAbsolutePath());
                    // 递归删除目录及其所有内容
                    Files.walk(topicDir.toPath())
                            .sorted(java.util.Comparator.reverseOrder())
                            .map(Path::toFile)
                            .forEach(File::delete);
                } catch (IOException e) {
                    log.error("Failed to delete partition directory: {}", topicDir.getAbsolutePath(), e);
                }
            }
        }
        log.info("Finished deleting storage files for topic '{}'.", topicName);
    }

    /**
     * 【RocketMQ风格】并行恢复所有分区，提升启动速度。
     * 应在Broker启动时调用。
     */
    public void recoverAllPartitions() {
        int threadCount = Math.max(2, Runtime.getRuntime().availableProcessors());
        ExecutorService pool = Executors.newFixedThreadPool(threadCount);
        List<Future<?>> futures = new ArrayList<>();
        for (LogPartition partition : sharedPartitions.values()) {
            futures.add(pool.submit(() -> {
                try {
                    partition.initializeAndRecover();
                } catch (IOException e) {
                    log.error("Failed to recover partition {}-{}", partition.getTopic(), partition.getPartitionId(), e);
                }
            }));
        }
        // 等待所有分区恢复完成
        for (Future<?> f : futures) {
            try { f.get(); } catch (Exception ignore) {}
        }
        pool.shutdown();
        log.info("所有分区并行恢复完成");
    }
}